import torch
from torch.utils.data import DataLoader
from data.dataset import EnergyTimeSeriesDataset
from models.model import EnhancedCNNLSTMAttention
from config.config import Config
from utils.metrics import calculate_metrics, print_metrics
import matplotlib.pyplot as plt
import numpy as np
import os
from sklearn.model_selection import train_test_split
import matplotlib as mpl
import matplotlib.font_manager as fm

def evaluate_model():
    config = Config()
    
    print("Loading dataset...")
    full_dataset = EnergyTimeSeriesDataset(config, config.DATA_PATH)
    print(f"Total dataset size: {len(full_dataset)}")
    
    # 数据集划分
    total_size = len(full_dataset)
    train_size = int(0.6 * total_size)
    val_size = int(0.2 * total_size)
    test_size = total_size - train_size - val_size
    print(f"Split sizes - Train: {train_size}, Val: {val_size}, Test: {test_size}")
    
    train_dataset, temp_dataset = torch.utils.data.random_split(
        full_dataset, 
        [train_size, val_size + test_size]
    )
    val_dataset, test_dataset = torch.utils.data.random_split(
        temp_dataset,
        [val_size, test_size]
    )
    
    # 创建数据加载器
    test_loader = DataLoader(test_dataset, batch_size=config.BATCH_SIZE, shuffle=False)
    val_loader = DataLoader(val_dataset, batch_size=config.BATCH_SIZE, shuffle=False)
    
    # 加载模型
    model = EnhancedCNNLSTMAttention(config).to(config.DEVICE)
    model_path = os.path.join(config.MODEL_SAVE_PATH, 'best_model.pth')
    
    if not os.path.exists(model_path):
        print(f"Error: Model file not found at {model_path}")
        print("Checking current directory for best_model.pth...")
        if os.path.exists('best_model.pth'):
            model_path = 'best_model.pth'
        else:
            raise FileNotFoundError("Could not find model file")
    
    print(f"Loading model from: {model_path}")
    checkpoint = torch.load(model_path)
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()
    
    # 评估验证集
    print("\nValidation Set Evaluation:")
    val_predictions, val_actuals = predict(model, val_loader, config)
    val_metrics = calculate_metrics(val_actuals, val_predictions)
    print_metrics(val_metrics)
    
    # 评估测试集
    print("\nTest Set Evaluation:")
    test_predictions, test_actuals = predict(model, test_loader, config)
    test_metrics = calculate_metrics(test_actuals, test_predictions)
    print_metrics(test_metrics)
    
    print(f"Saving plots to results/ directory...")
    plot_all_features(test_predictions, test_actuals, config)
    print("Evaluation completed!")

def predict(model, loader, config):
    predictions = []
    actuals = []
    
    with torch.no_grad():
        for batch_x, batch_y in loader:
            batch_x = batch_x.to(config.DEVICE)
            outputs = model(batch_x)
            predictions.extend(outputs.cpu().numpy())
            actuals.extend(batch_y.numpy())
    
    return np.array(predictions), np.array(actuals)

def plot_all_features(predictions, actuals, config, save_path='results/'):
    # macOS 系统字体设置
    import matplotlib.font_manager as fm
    
    # 查找系统字体
    font_paths = [
        '/System/Library/Fonts/PingFang.ttc',  # macOS 默认中文字体
        '/System/Library/Fonts/STHeiti Light.ttc',  # macOS 另一个中文字体
        '/System/Library/Fonts/STHeiti Medium.ttc'
    ]
    
    font_found = False
    for font_path in font_paths:
        if os.path.exists(font_path):
            prop = fm.FontProperties(fname=font_path)
            plt.rcParams['font.family'] = prop.get_name()
            font_found = True
            print(f"Using font from: {font_path}")
            break
    
    if not font_found:
        print("Warning: Using system default font...")
        plt.rcParams['font.sans-serif'] = ['Arial']
        
    plt.rcParams['axes.unicode_minus'] = False
    
    os.makedirs(save_path, exist_ok=True)
    
    for i, feature_name in enumerate(config.FEATURE_COLUMNS):
        plt.figure(figsize=(15, 6))
        plt.plot(actuals[:, i], label='实际值', alpha=0.7)
        plt.plot(predictions[:, i], label='预测值', alpha=0.7)
        
        # 直接使用特征名（已经是中文）
        plt.title(f'{feature_name}的预测结果', fontsize=12, fontproperties=prop if font_found else None)
        plt.xlabel('时间步', fontsize=10, fontproperties=prop if font_found else None)
        plt.ylabel('数值', fontsize=10, fontproperties=prop if font_found else None)
        plt.legend(prop=prop if font_found else None, fontsize=10)
        plt.grid(True)
        
        # 添加评估指标到图表
        metrics = calculate_metrics(actuals[:, i], predictions[:, i])
        plt.text(0.02, 0.98, 
                f'均方误差: {metrics["MSE"]:.4f}\n'
                f'R²值: {metrics["R2"]:.4f}\n'
                f'平均绝对百分比误差: {metrics["MAPE"]:.2f}%', 
                transform=plt.gca().transAxes, 
                verticalalignment='top',
                bbox=dict(boxstyle='round', facecolor='white', alpha=0.8),
                fontsize=10,
                fontproperties=prop if font_found else None)
        
        # 保存图片
        plt.savefig(os.path.join(save_path, f'{feature_name}_预测结果.png'), 
                   bbox_inches='tight', 
                   dpi=300)
        plt.close()

if __name__ == '__main__':
    evaluate_model() 